#!/usr/bin/env python3
"""
自动化题目搬运脚本，包含PDF转图片、答案处理和题面文件生成功能
"""

import os
import shutil
from pdf2image import convert_from_path
import tkinter as tk
from tkinter import filedialog
import re
import sys


def convert_pdf_to_images(pdf_path=None):
    """将PDF文件的每一页转换为图片并保存到additional_file目录中"""
    # 如果没有提供PDF路径，则使用tkinter文件对话框
    if not pdf_path:
        root = tk.Tk()
        root.withdraw()
        pdf_path = filedialog.askopenfilename(
            title="选择PDF文件",
            filetypes=[("PDF files", "*.pdf")]
        )
    
    # 棜查是否选择了文件
    if not pdf_path:
        print("未选择PDF文件，程序退出。")
        return None, 0
    
    # 获取PDF文件名（不含扩展名），并去除空格
    pdf_filename = os.path.splitext(os.path.basename(pdf_path))[0]
    pdf_filename = pdf_filename.replace(" ", "_")
    
    # 获取脚本所在目录的绝对路径
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 目标目录（使用绝对路径）
    output_dir = os.path.join(script_dir, "additional_file")
    
    # 删除目标目录下的所有文件
    if os.path.exists(output_dir):
        for filename in os.listdir(output_dir):
            file_path = os.path.join(output_dir, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f'删除 {file_path} 失败. 原因: {e}')
    
    # 确保目标目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 将PDF转换为图片
    print("正在转换PDF为图片...")
    pages = convert_from_path(pdf_path, dpi=200)
    
    # 保存每一页为图片，文件名包含原PDF文件名
    for i, page in enumerate(pages):
        image_path = os.path.join(output_dir, f"{pdf_filename}_page_{i+1:02d}.png")
        page.save(image_path, "PNG")
        print(f"已保存: {image_path}")
    
    print(f"PDF转图片完成，共保存了 {len(pages)} 张图片。")
    return pdf_filename, len(pages)


def process_answers():
    """处理答案文件，生成测试数据"""
    # 获取脚本所在目录的绝对路径
    path = os.path.dirname(os.path.abspath(__file__))
    
    # 创建测试数据目录
    testdata_dir = os.path.join(path, "testdata")
    if not os.path.isdir(testdata_dir):
        os.mkdir(testdata_dir)
    
    # 读取答案文件
    with open(os.path.join(path, "answer.txt"), "r", encoding="utf-8") as f:
        ans = f.read()
    
    # 处理答案字符串
    newans = ""
    keys = ["√", "×", "x", "X"]
    for i in range(len(ans)):
        if ans[i] in keys:
            newans += ans[i]
        elif ans[i] >= "A" and ans[i] <= "Z":
            newans += ans[i]
    
    ans = newans
    ans = ans.replace(" ", "")
    ans = ans.replace("\r", "")
    ans = ans.replace("\n", "")
    ans = ans.replace("\t", "")
    ans = ans.replace("√", "T")
    ans = ans.replace("×", "F")
    ans = ans.replace("x", "F")
    ans = ans.replace("X", "F")
    
    # 生成测试数据文件
    for i in range(1, len(ans)+1):
        with open(os.path.join(testdata_dir, f"{i}.in"), "w", encoding="utf-8") as f:
            f.write(str(i))
        with open(os.path.join(testdata_dir, f"{i}.out"), "w", encoding="utf-8") as f:
            f.write(ans[i-1])
    
    print(f"答案处理完成，共生成 {len(ans)} 个测试数据文件。")
    print("答案序列:", ans)
    return len(ans)


def generate_problem_md(pdf_filename, page_count):
    """生成题面文件problem_zh.md"""
    # 获取脚本所在目录的绝对路径
    path = os.path.dirname(os.path.abspath(__file__))
    problem_file = os.path.join(path, "problem_zh.md")
    
    # 如果已存在题面文件，则先备份
    if os.path.exists(problem_file):
        shutil.copy2(problem_file, problem_file + ".bak")
        print(f"已备份原题面文件为: {problem_file}.bak")
    
    # 生成新的题面文件
    with open(problem_file, "w", encoding="utf-8") as f:
        
        f.write("<center>\n\n# " + pdf_filename + "\n</center>\n\n---\n\n")
        
        
        f.write("# 一、客观题\n")
        # 写入每页图片
        for i in range(1, page_count + 1):
            f.write(f"![alt text](file://{pdf_filename}_page_{i:02d}.png)\n\n")

        # 添加编程题链接占位符
        f.write("# 二、编程题\n")
        for i in range(1, 2):  # 为第两道编程题添加占位符
            
            f.write("<!-- 事后补完编程题再填上连接 -->\n")
            f.write(f"### **[编程题{i}](https://fslong.iok.la/problem/GESP20250601T{i})**\n")
    
    print(f"题面文件生成完成: {problem_file}")


def generate_problem_yaml(pdf_filename):
    """生成题目元数据文件problem.yaml"""
    # 获取脚本所在目录的绝对路径
    path = os.path.dirname(os.path.abspath(__file__))
    yaml_file = os.path.join(path, "problem.yaml")
    
    # 中文级别到数字的映射
    level_map = {
        '一': '1',
        '二': '2',
        '三': '3',
        '四': '4',
        '五': '5',
        '六': '6',
        '七': '7',
        '八': '8',
        '九': '9',
        '十': '10'
    }
    
    # 从文件名中提取日期和级别信息
    # 匹配类似 202506-GESP-C++-四级 的格式
    pattern = r'(\d{6}).*?[Gg][Ee][Ss][Pp].*?[Cc\+]{2,4}.*?([一二三四五六七八九十\d])[级級]'
    match = re.search(pattern, pdf_filename)
    
    if match:
        date = match.group(1)  # 日期，如 202506
        level_raw = match.group(2)  # 级别，如 四 或 4
        # 如果是中文级别，则转换为数字
        level = level_map.get(level_raw, level_raw)
        pid = f"gesp{date}cpp{level}"
    else:
        # 如果无法匹配，则使用默认格式
        pid = "gesp" + re.sub(r'[^\w]', '', pdf_filename.lower())
        level = "unknown"
    
    # 如果已存在元数据文件，则先备份
    if os.path.exists(yaml_file):
        shutil.copy2(yaml_file, yaml_file + ".bak")
        print(f"已备份原元数据文件为: {yaml_file}.bak")
    
    # 生成新的元数据文件
    with open(yaml_file, "w", encoding="utf-8") as f:
        f.write(f"pid: {pid}\n")
        f.write(f'title: "{pdf_filename}(客观题)"\n')
        f.write("tag: \n")
        f.write("  - GESP\n")
        if level != "unknown":
            f.write(f"  - GESP{level}级\n")
    
    print(f"题目元数据文件生成完成: {yaml_file}")


def main():
    """主函数，执行完整的题目搬运自动化流程"""
    # 检查命令行参数
    if "-A" in sys.argv:
        # 只执行答案处理功能
        print("开始执行答案处理功能...")
        try:
            answer_count = process_answers()
            print("答案处理完成！")
        except FileNotFoundError:
            print("错误：找不到answer.txt文件，请确认文件是否存在。")
        except Exception as e:
            print(f"处理答案时出错：{e}")
        return
    
    print("开始执行题目搬运自动化流程...")
    
    # 检查是否通过命令行参数提供了PDF文件路径
    pdf_path = None
    if len(sys.argv) > 1 and sys.argv[1] != "-A":
        pdf_path = sys.argv[1]
        if not os.path.exists(pdf_path):
            print(f"指定的PDF文件不存在: {pdf_path}")
            return
    
    # 转换PDF为图片
    pdf_filename, page_count = convert_pdf_to_images(pdf_path)
    
    # 检查是否成功选择了PDF文件
    if pdf_filename is None:
        print("未选择PDF文件，程序退出。")
        return
    
    # 处理答案
    try:
        answer_count = process_answers()    
    except FileNotFoundError:
        print("错误：找不到answer.txt文件，请确认文件是否存在。")
        return
    except Exception as e:
        print(f"处理答案时出错：{e}")
        return
        
    pdf_filename = pdf_filename.replace(" ", "_")
    # 生成题面文件
    generate_problem_md(pdf_filename, page_count)
    
    # 生成题目元数据文件
    generate_problem_yaml(pdf_filename)
    
    print("题目搬运自动化流程执行完成！")


if __name__ == "__main__":
    main()